#ifndef ECHO_NET
#define ECHO_NET

#include "const.h"
#include "matrix.h"
#include "echo_input_weights.h"
#include "echo_internal_weights.h"

using namespace std;

class EchoNet
{
	private:
		Matrix<float> input_weights;
		Matrix<float> internal_weights;
		Matrix<float> rand_noise;
		int n_internal_units;
		float input_scaling;
		float noise_level;
		float leak;

	public:
		EchoNet();
		
		Matrix<float> getState(Matrix<float>&);
};

EchoNet::EchoNet() : input_weights(N_INTERNAL_UNITS, V), internal_weights(N_INTERNAL_UNITS, N_INTERNAL_UNITS), rand_noise(N_INTERNAL_UNITS, 1)
{
	if (IF_LOAD_MODEL)
	{
		input_weights.setValue(_input_weights);
		internal_weights.setValue(_internal_weights);
		n_internal_units = N_INTERNAL_UNITS;
		input_scaling 	 = INPUT_SCALING;
		noise_level 	 = NOISE_LEVEL;
		leak 			 = LEAK;
	}
}

Matrix<float> EchoNet::getState(Matrix<float>& inputData)
{
	int timeSteps = T;
	int dataDims = V;
	Matrix<float> stateMatrix(timeSteps, N_INTERNAL_UNITS);
	if (inputData.rows() != timeSteps || inputData.cols() != dataDims) {
		return stateMatrix;
	}
	Matrix<float> preState(1, N_INTERNAL_UNITS);
	Matrix<float> preState_T(N_INTERNAL_UNITS, 1);
	Matrix<float> curInput(1, dataDims);
	Matrix<float> stateBeforeTanh(N_INTERNAL_UNITS, 1);
	preState.setZeroMatrix();
	preState_T.setZeroMatrix();
	stateMatrix.setZeroMatrix();
	curInput.setZeroMatrix();

	for (int i = 0; i < timeSteps; i++) {
		for (int j = 0; j < dataDims; j++) {
			curInput(0, j) = inputData(i, j);
		}
		
		stateBeforeTanh = internal_weights * preState_T + input_weights * curInput.transpose();
		rand_noise.setRandMatrix();
		stateBeforeTanh = stateBeforeTanh + rand_noise * noise_level;
		
		Matrix<float> curState(1, N_INTERNAL_UNITS);
		curState = stateBeforeTanh.tanhMatrix();
		if (leak <= 0.000001) {	
			preState = curState.transpose();
		}
		else {
			preState = preState * (1.0 - leak) + curState;
		}

		for (int j = 0; j < N_INTERNAL_UNITS; j++) {
			stateMatrix(i, j) = preState(0, j);
		}
	}

	return stateMatrix;
}

#endif